import shutil

import cv2
import insightface
import numpy as np
import os


class FaceEntity:
    def __init__(self, id, image, embedding, confidence, count):
        self.id = id
        self.image = image
        self.embedding = embedding
        self.confidence = confidence
        self.count = count


def main(video_path, output_path):
    # 加载人脸识别模型
    model = insightface.app.FaceAnalysis()
    model.prepare(ctx_id=-1, nms=0.4)

    # 创建保存人脸实体的列表
    face_entities = []

    file_path, file_name = os.path.split(video_path)
    name, ext = os.path.splitext(file_name)
    output_path = os.path.join(output_path, name)
    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    os.mkdir(output_path)
    # 创建视频捕获对象
    cap = cv2.VideoCapture(video_path)

    # 检查视频是否打开成功
    if not cap.isOpened():
        print("Error: Could not open video.")
        exit()

    # 定义帧计数器
    frame_count = 0

    # 循环读取视频每一帧
    while True:
        # 读取下一帧
        ret, frame = cap.read()

        # 如果正确读取帧，ret为True
        if not ret:
            print("Done: No more frames.")
            break

        # 保存帧到文件
        img_file = os.path.join(output_path, frame_count + ".png")
        cv2.imwrite(img_file, frame)

        # 进行人脸识别和定位
        faces = model.get(frame)

        # 遍历识别出的人脸
        for idx, face in enumerate(faces):
            # 获取人脸框的位置
            bbox = face.bbox.astype(np.int)

            # 裁剪人脸图像
            face_image = frame[bbox[1]:bbox[3], bbox[0]:bbox[2]]

            # 获取人脸特征数据
            face_embedding = face.embedding

            # 获取人脸置信度
            confidence = face.normed_embedding.norm()

            # 判断是否存在相同人脸
            exist, index = is_face_exist(face_entities, face_embedding)

            if exist:
                # 如果存在相同人脸，则增加检测次数
                face_entity = face_entities[index]
                face_entity.count += 1
                if face_entity.confidence < confidence:
                    face_entity.image = img_file
            else:
                # 创建人脸实体对象
                face_entity = FaceEntity(idx, img_file, face_embedding, confidence, 1)

                # 将人脸实体对象添加到列表中
                face_entities.append(face_entity)

        print("人脸识别完成。")
        frame_count += 1

    # 释放视频捕获对象
    cap.release()

    # 遍历人脸实体列表，输出人脸信息
    for face_entity in face_entities:
        print("人脸ID:", face_entity.id)
        print("人脸置信度:", face_entity.confidence)
        print("人脸检测次数:", face_entity.count)
        # 可以根据需要保存人脸图像和特征数据到本地或进行其他处理


def is_face_exist(face_entities, face_embedding):
    # 遍历人脸实体列表，判断是否存在相同人脸
    for idx, face_entity in enumerate(face_entities):
        # 计算人脸特征之间的距离
        distance = np.linalg.norm(face_entity.embedding - face_embedding)

        # 如果距离小于阈值，则认为是相同人脸
        if distance < 0.3:
            return True, idx

    return False, -1


if __name__ == "__main__":
    video_file = ""
    main(video_file, "")
